/*
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov.example.chain;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.io.StreamCorruptedException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Map.Entry;

import org.apache.commons.math3.stat.descriptive.SummaryStatistics;

import net.javlov.Action;
import net.javlov.DecayingLearningRate;
import net.javlov.EpisodicRewardStepStatistic;
import net.javlov.FixedLearningRate;
import net.javlov.Option;
import net.javlov.Policy;
import net.javlov.QLearningAgent;
import net.javlov.SarsaAgent;
import net.javlov.Simulator;
import net.javlov.State;
import net.javlov.TDAgent;
import net.javlov.policy.EGreedyPolicy;
import net.javlov.policy.RandomPolicy;
import net.javlov.vf.*;

@SuppressWarnings("all")
public class MDPMain implements Runnable {

	Simulator sim;
	QMap<Integer> vt;
	
	public static void main(String[] args) {
		MDPMain m = new MDPMain();
		m.init();
		m.run();
	}
	
	public void init() {
		//params
		double gamma = 1, lambda = 0.9, epsilon = 0.1;
		
		//world
		Chain world = new Chain(10);
		
		//actions
		List optionPool = new ArrayList(2);
		Move left = new Move(world, Move.LEFT),
			right = new Move(world, Move.RIGHT);
		optionPool.add(left);
		optionPool.add(right);
		
		//value function
		vt = new QMap<Integer>(optionPool, 15); //for TD(0)
		OptionTraces trs = new OptionTraces(new BasicOptionTrace(gamma, lambda,1,false), 0.001);
		TracedQTable<Integer> tqt = new TracedQTable<Integer>( vt, trs);
		
		//policy
		Policy pi = new EGreedyPolicy(vt, epsilon, optionPool);
		
		//agent
		SarsaAgent agent = new QLearningAgent(vt, gamma);
		agent.setPolicy(pi);
		agent.setLearnRate( new DecayingLearningRate(1,2,0.8) );
		//agent.setLearnRate( new FixedLearningRate(0.2) );
		
		//simulator / glue
		sim = new Simulator();
		sim.setAgent(agent);
		sim.setEnvironment(world);		
	}
	
	public void run() {
		//System.out.println("START");
		//QMap<Integer> convMap = (QMap)loadObject("convergedqtable.ser");
		int episodes = 150;
		EpisodicRewardStepStatistic stat = new EpisodicRewardStepStatistic(episodes);
		SummaryStatistics sstat[] = new SummaryStatistics[episodes];
		for ( int i = 0; i < episodes; i++ ) {
			sstat[i] = new SummaryStatistics();
		}
		sim.addStatistic(stat);
		
		for ( int r = 0; r < 100; r++ ) {
			sim.init();
			for ( int i = 0; i < episodes; i++ ) {
				sim.runEpisode();
				sim.reset();
				//System.out.println("================");
				//System.out.println(calcTableError(convMap, vt));
				sstat[i].addValue( stat.getSteps()[i] );
				//sstat[i].addValue( stat.getRewards()[i] );
			}
		}
		for ( int i = 0; i < episodes; i++ ) {
			System.out.println( sstat[i].getMean() );
		}
		//serializeObject(vt, "convergedtable.ser");
	}
	
	public <T> double calcTableError(ValueMap<T> convMap, ValueMap<T> newMap) {
		double err = 0;
		for ( Map.Entry<State<? extends T>, Value> svPair : convMap.stateValueSet() ) {
			err += Math.pow(svPair.getValue().get() - newMap.getValue(svPair.getKey()),2);
		}
		return err;
	}
	
	public static void serializeObject(Serializable obj, String fileName) {
        ObjectOutputStream oos = null;
        try {
                oos = new ObjectOutputStream( new FileOutputStream(fileName) );
                oos.writeObject(obj);
        } catch (IOException e) {
                System.out.println("Oops, an IOException occurred: " + e);
        } finally {
                try {
                        oos.close();
                } catch (IOException e) {
                        System.out.println("Oops, an IOException occurred: " + e);
                }
        }
	}

	public static Object loadObject(String filename) {
        Object obj = null;
        ObjectInputStream ois = null;
        try {
                ois = new ObjectInputStream( new FileInputStream(filename) );
                obj = ois.readObject();
        } catch (StreamCorruptedException sce) {
                System.err.println("Not a binary file: " + filename);
        } catch (Exception e) {
                System.err.println("WARNING: Problem trying to restore object from file: " + filename 
                                + ": " + e + ".");
        } finally {
                try {
                        if ( ois != null )
                                ois.close();
                } catch (IOException ioe) {
                        System.err.println("WARNING: problem closing stream: " + ioe);
                }
        }
        
        return obj;
	}
}
